import numpy as np
import pandas as pd
from scipy import stats, signal
import pywt
from typing import List, Dict, Tuple
import os
from tqdm import tqdm


class BearingFeatureExtractor:
    def __init__(self, sampling_frequency: int = 25600):
        """
        初始化轴承特征提取器

        参数:
            sampling_frequency: 采样频率(Hz)
        """
        self.fs = sampling_frequency
        # 中文特征名称
        self.time_domain_feature_names = [
            '均值', '标准差', '均方根值', '最大值', '最小值', '峰峰值',
            '偏度', '峰度', '峰值因子', '裕度因子', '波形因子', '脉冲因子',
            '绝对平均值', '方根幅值', '平均幅值', '短时能量', '过零率'
        ]

        # 频带分布为每5Hz一个区间
        self.freq_bands = [(i, i + 5) for i in range(0, 100, 5)]  # 0-100Hz，每5Hz一个区间
        self.freq_domain_feature_names = [
                                             '频域均值', '频域标准差', '频域均方根值', '频域最大值',
                                             '频域最小值', '频域峰值', '频域方差', '频域偏度',
                                             '频域峰度', '频谱质心', '频谱扩展', '频谱熵', '频谱滚降',
                                             '重心频率', '均方频率', '均方根频率', '频率方差', '频率标准差',
                                             '基频', '峰值频率'
                                         ] + [f'频带{i}-{i + 5}Hz能量' for i in range(0, 100, 5)]

        self.time_freq_feature_names = [
            '小波能量熵', '小波奇异熵', 'STFT能量均值', 'STFT能量标准差',
            'STFT能量熵', 'STFT谱质心', 'STFT谱扩展'
        ]

        # 小波分解特征 - 计算1-5层所有细节系数
        self.wavelet_feature_names = []
        for level in range(1, 6):
            self.wavelet_feature_names.extend([
                f'level{level}d_abs_mean', f'level{level}d_energy',
                f'level{level}d_energy_smax', f'level{level}d_energy_sstd'
            ])

        self.wavelet_band_names = [f'db3_{i}频带能量比' for i in range(8)]

    def extract_all_features(self, signal_data: np.ndarray) -> Dict[str, float]:
        """
        提取信号的所有特征（时域、频域、时频域）

        参数:
            signal_data: 输入信号数据

        返回:
            包含所有特征的字典
        """
        # 检查信号是否有效
        if len(signal_data) == 0:
            return {name: 0 for name in
                    self.time_domain_feature_names +
                    self.freq_domain_feature_names +
                    self.time_freq_feature_names +
                    self.wavelet_feature_names +
                    self.wavelet_band_names}

        # 提取时域特征
        time_features = self.extract_time_domain_features(signal_data)

        # 提取频域特征
        freq_features = self.extract_frequency_domain_features(signal_data)

        # 提取时频域特征
        time_freq_features = self.extract_time_frequency_features(signal_data)

        # 提取小波分解特征
        wavelet_features = self.extract_wavelet_features(signal_data)

        # 合并所有特征
        all_features = {**time_features, **freq_features, **time_freq_features, **wavelet_features}
        return all_features

    def extract_time_domain_features(self, signal_data: np.ndarray) -> Dict[str, float]:
        """
        提取时域特征

        参数:
            signal_data: 输入信号数据

        返回:
            时域特征字典
        """
        features = {}

        # 基本统计特征
        features['均值'] = np.mean(signal_data)
        features['标准差'] = np.std(signal_data)
        features['最大值'] = np.max(signal_data)
        features['最小值'] = np.min(signal_data)

        # 均方根值
        features['均方根值'] = np.sqrt(np.mean(signal_data ** 2))

        # 峰峰值
        features['峰峰值'] = features['最大值'] - features['最小值']

        # 偏度和峰度
        features['偏度'] = stats.skew(signal_data)
        features['峰度'] = stats.kurtosis(signal_data)

        # 波形指标
        features['峰值因子'] = features['最大值'] / features['均方根值'] if features['均方根值'] != 0 else 0
        features['裕度因子'] = features['最大值'] / (np.mean(np.sqrt(np.abs(signal_data)))) ** 2 if np.mean(
            np.sqrt(np.abs(signal_data))) != 0 else 0
        features['波形因子'] = features['均方根值'] / np.mean(np.abs(signal_data)) if np.mean(
            np.abs(signal_data)) != 0 else 0
        features['脉冲因子'] = features['最大值'] / np.mean(np.abs(signal_data)) if np.mean(
            np.abs(signal_data)) != 0 else 0

        # 新增时域特征
        features['绝对平均值'] = np.mean(np.abs(signal_data))
        features['方根幅值'] = (np.mean(np.sqrt(np.abs(signal_data)))) ** 2
        features['平均幅值'] = np.mean(np.abs(signal_data))

        # 短时能量
        features['短时能量'] = np.sum(signal_data ** 2) / len(signal_data)

        # 过零率
        zero_crossings = np.where(np.diff(np.sign(signal_data)))[0]
        features['过零率'] = len(zero_crossings) / len(signal_data) if len(signal_data) > 0 else 0

        return features

    def extract_frequency_domain_features(self, signal_data: np.ndarray) -> Dict[str, float]:
        """
        提取频域特征

        参数:
            signal_data: 输入信号数据

        返回:
            频域特征字典
        """
        # 计算FFT
        n = len(signal_data)
        if n < 2:  # 至少需要两个点进行FFT
            return {name: 0 for name in self.freq_domain_feature_names}

        fft_data = np.fft.fft(signal_data)
        fft_magnitude = np.abs(fft_data[:n // 2])
        fft_freq = np.fft.fftfreq(n, 1 / self.fs)[:n // 2]

        # 功率谱密度
        psd = (2.0 / n) * fft_magnitude ** 2

        features = {}

        # 基本统计特征
        features['频域均值'] = np.mean(psd)
        features['频域标准差'] = np.std(psd)
        features['频域均方根值'] = np.sqrt(np.mean(psd ** 2))
        features['频域最大值'] = np.max(psd)
        features['频域最小值'] = np.min(psd)
        features['频域峰值'] = np.max(psd)
        features['频域方差'] = np.var(psd)
        features['频域偏度'] = stats.skew(psd)
        features['频域峰度'] = stats.kurtosis(psd)

        # 频谱质心
        if np.sum(psd) > 0:
            features['频谱质心'] = np.sum(fft_freq * psd) / np.sum(psd)
        else:
            features['频谱质心'] = 0

        # 频谱扩展
        if features['频谱质心'] != 0 and np.sum(psd) > 0:
            features['频谱扩展'] = np.sqrt(np.sum(((fft_freq - features['频谱质心']) ** 2) * psd) / np.sum(psd))
        else:
            features['频谱扩展'] = 0

        # 频谱熵
        normalized_psd = psd / np.sum(psd) if np.sum(psd) > 0 else psd
        normalized_psd = normalized_psd[normalized_psd > 0]  # 避免log(0)
        features['频谱熵'] = -np.sum(normalized_psd * np.log2(normalized_psd)) if len(normalized_psd) > 0 else 0

        # 频谱滚降点(95%)
        cumsum_psd = np.cumsum(psd)
        if np.sum(psd) > 0:
            features['频谱滚降'] = fft_freq[np.where(cumsum_psd >= 0.95 * cumsum_psd[-1])[0][0]]
        else:
            features['频谱滚降'] = 0

        # 新增频域特征
        # 重心频率
        features['重心频率'] = np.sum(fft_freq * psd) / np.sum(psd) if np.sum(psd) > 0 else 0

        # 均方频率
        features['均方频率'] = np.sum((fft_freq ** 2) * psd) / np.sum(psd) if np.sum(psd) > 0 else 0

        # 均方根频率
        features['均方根频率'] = np.sqrt(features['均方频率']) if features['均方频率'] >= 0 else 0

        # 频率方差
        if features['重心频率'] != 0 and np.sum(psd) > 0:
            features['频率方差'] = np.sum(((fft_freq - features['重心频率']) ** 2) * psd) / np.sum(psd)
        else:
            features['频率方差'] = 0

        # 频率标准差
        features['频率标准差'] = np.sqrt(features['频率方差']) if features['频率方差'] >= 0 else 0

        # 基频(使用自相关法估计)
        autocorr = np.correlate(signal_data, signal_data, mode='full')
        autocorr = autocorr[len(autocorr) // 2:]
        peak_indices = signal.find_peaks(autocorr)[0]
        if len(peak_indices) > 1:
            features['基频'] = self.fs / peak_indices[1]  # 第一个峰值后的最大峰值
        else:
            features['基频'] = 0

        # 峰值频率
        max_psd_idx = np.argmax(psd)
        features['峰值频率'] = fft_freq[max_psd_idx]

        # 计算频带能量(0-100Hz，每5Hz一个区间)
        total_energy = np.sum(psd)
        for i, (f_start, f_end) in enumerate(self.freq_bands):
            band_mask = (fft_freq >= f_start) & (fft_freq < f_end)
            band_energy = np.sum(psd[band_mask])
            features[f'频带{f_start}-{f_end}Hz能量'] = band_energy / total_energy if total_energy > 0 else 0

        return features

    def extract_time_frequency_features(self, signal_data: np.ndarray) -> Dict[str, float]:
        """
        提取时频域特征

        参数:
            signal_data: 输入信号数据

        返回:
            时频域特征字典
        """
        features = {}

        # 确保信号长度足够进行STFT
        if len(signal_data) < 256:  # STFT的窗口大小
            return {name: 0 for name in self.time_freq_feature_names}

        # 1. 短时傅里叶变换(STFT)
        f, t, Zxx = signal.stft(signal_data, fs=self.fs, nperseg=256, noverlap=128)
        stft_mag = np.abs(Zxx)

        # STFT能量特征
        stft_energy = np.sum(stft_mag ** 2, axis=0)
        features['STFT能量均值'] = np.mean(stft_energy)
        features['STFT能量标准差'] = np.std(stft_energy)

        # STFT能量熵
        normalized_stft_energy = stft_energy / np.sum(stft_energy) if np.sum(stft_energy) > 0 else stft_energy
        normalized_stft_energy = normalized_stft_energy[normalized_stft_energy > 0]
        features['STFT能量熵'] = -np.sum(normalized_stft_energy * np.log2(normalized_stft_energy)) if len(
            normalized_stft_energy) > 0 else 0

        # STFT谱特征
        time_centers = np.arange(len(stft_energy))
        if np.sum(stft_energy) > 0:
            features['STFT谱质心'] = np.sum(time_centers * stft_energy) / np.sum(stft_energy)
        else:
            features['STFT谱质心'] = 0

        if features['STFT谱质心'] != 0 and np.sum(stft_energy) > 0:
            features['STFT谱扩展'] = np.sqrt(
                np.sum(((time_centers - features['STFT谱质心']) ** 2) * stft_energy) / np.sum(stft_energy))
        else:
            features['STFT谱扩展'] = 0

        # 2. 连续小波变换(CWT)
        # 使用morlet小波
        widths = np.arange(1, 31)
        try:
            cwtmatr = signal.cwt(signal_data, signal.morlet2, widths)
            cwt_energy = np.sum(cwtmatr ** 2, axis=0)

            # 小波能量熵
            normalized_cwt_energy = cwt_energy / np.sum(cwt_energy) if np.sum(cwt_energy) > 0 else cwt_energy
            normalized_cwt_energy = normalized_cwt_energy[normalized_cwt_energy > 0]
            features['小波能量熵'] = -np.sum(normalized_cwt_energy * np.log2(normalized_cwt_energy)) if len(
                normalized_cwt_energy) > 0 else 0

            # 小波奇异熵(使用SVD分解)
            U, s, V = np.linalg.svd(cwtmatr, full_matrices=False)
            normalized_s = s / np.sum(s) if np.sum(s) > 0 else s
            features['小波奇异熵'] = -np.sum(normalized_s * np.log2(normalized_s)) if len(normalized_s) > 0 else 0
        except Exception as e:
            print(f"小波变换出错: {e}")
            features['小波能量熵'] = 0
            features['小波奇异熵'] = 0

        return features

    def extract_wavelet_features(self, signal_data: np.ndarray) -> Dict[str, float]:
        """
        提取小波分解特征（计算1-5层所有细节系数特征）

        参数:
            signal_data: 输入信号数据

        返回:
            小波分解特征字典
        """
        features = {}

        # 确保信号长度足够进行5层小波分解
        min_length = 2 ** 5  # 5层分解所需的最小长度
        if len(signal_data) < min_length:
            return {name: 0 for name in self.wavelet_feature_names + self.wavelet_band_names}

        try:
            # 进行5层小波分解
            coeffs = pywt.wavedec(signal_data, 'db4', level=5)

            # 计算每一层细节系数特征
            for level in range(1, 6):
                # 获取当前层的细节系数
                d_coeffs = coeffs[6 - level]  # 注意coeffs的索引顺序

                # 计算基本特征
                features[f'level{level}d_abs_mean'] = np.mean(np.abs(d_coeffs))
                features[f'level{level}d_energy'] = np.sum(d_coeffs ** 2)

                # 将细节系数分帧计算能量统计特征
                frame_size = min(32, len(d_coeffs) // 2)  # 确保至少有两帧
                frames = [d_coeffs[i:i + frame_size] for i in range(0, len(d_coeffs) - frame_size + 1, frame_size)]
                frame_energies = [np.sum(frame ** 2) for frame in frames]

                if frame_energies:
                    features[f'level{level}d_energy_smax'] = np.max(frame_energies)
                    features[f'level{level}d_energy_sstd'] = np.std(frame_energies)
                else:
                    features[f'level{level}d_energy_smax'] = 0
                    features[f'level{level}d_energy_sstd'] = 0

            # 小波包分解特征
            wp = pywt.WaveletPacket(data=signal_data, wavelet='db3', mode='symmetric', maxlevel=3)
            nodes = [node.path for node in wp.get_level(3, 'natural')]
            energy = [np.sum(wp[node].data ** 2) for node in nodes]
            total_energy = np.sum(energy)

            if total_energy > 0:
                energy_ratios = [e / total_energy for e in energy]
                for i in range(8):
                    features[f'db3_{i}频带能量比'] = energy_ratios[i]
            else:
                for i in range(8):
                    features[f'db3_{i}频带能量比'] = 0
        except Exception as e:
            print(f"小波分解出错: {e}")
            # 返回默认值
            for name in self.wavelet_feature_names + self.wavelet_band_names:
                features[name] = 0

        return features


def load_single_csv_file(file_path: str) -> Tuple[np.ndarray, np.ndarray]:
    """
    加载单个CSV文件，提取水平和竖直方向加速度数据

    参数:
        file_path: CSV文件路径

    返回:
        元组(水平加速度数据, 竖直加速度数据)
    """
    try:
        # 使用pandas读取CSV文件
        df = pd.read_csv(file_path, header=None)

        # 检查并转换数据类型
        df = df.apply(pd.to_numeric, errors='coerce')  # 将非数值转换为NaN

        # 处理缺失值
        df = df.fillna(0)  # 用0填充缺失值，也可以选择其他填充方法

        # 提取第一列(水平加速度)和第二列(竖直加速度)
        horizontal = df.iloc[:, 0].values
        vertical = df.iloc[:, 1].values

        # 检查数据类型
        if horizontal.dtype != np.float64:
            horizontal = horizontal.astype(np.float64)
        if vertical.dtype != np.float64:
            vertical = vertical.astype(np.float64)

        print(f"成功加载数据: 水平方向数据点 {len(horizontal)}, 竖直方向数据点 {len(vertical)}")
        print(f"数据类型: {horizontal.dtype}")  # 打印数据类型用于调试
        return horizontal, vertical
    except Exception as e:
        print(f"加载数据时出错: {e}")
        return np.array([]), np.array([])


def process_single_file(file_path: str, sampling_freq: int = 25600,
                        window_sec: float = 1.28, overlap_ratio: float = 0.0) -> pd.DataFrame:
    """
    处理单个CSV文件并返回特征DataFrame

    参数:
        file_path: 输入CSV文件路径
        window_sec: 窗口大小(秒)
        overlap_ratio: 重叠率(0-1之间)
    """
    # 加载数据
    horizontal, vertical = load_single_csv_file(file_path)

    if horizontal.size == 0 or vertical.size == 0:
        print(f"数据加载失败: {file_path}")
        return pd.DataFrame()

    # 将窗口大小(秒)转换为数据点数
    window_size = int(sampling_freq * window_sec)
    print(f"处理文件: {os.path.basename(file_path)}, 窗口大小: {window_sec}秒 = {window_size}数据点")

    # 初始化特征提取器
    extractor = BearingFeatureExtractor(sampling_freq)

    # 存储所有样本的特征
    all_features = []

    # 处理水平方向数据
    step = max(1, int(window_size * (1 - overlap_ratio)))
    samples = [(horizontal, vertical, 0)]  # 整个文件作为一个样本

    # 提取每个样本的特征
    for h_sample, v_sample, idx in samples:
        # 提取水平方向特征
        h_features = extractor.extract_all_features(h_sample)
        h_features['方向'] = '水平'
        h_features['文件编号'] = int(os.path.basename(file_path).split('.')[0])  # 从文件名提取编号
        h_features['样本编号'] = idx
        h_features['时间点'] = idx / sampling_freq

        # 提取竖直方向特征
        v_features = extractor.extract_all_features(v_sample)
        v_features['方向'] = '竖直'
        v_features['文件编号'] = int(os.path.basename(file_path).split('.')[0])
        v_features['样本编号'] = idx
        v_features['时间点'] = idx / sampling_freq

        all_features.extend([h_features, v_features])

    # 转换为DataFrame
    if all_features:
        features_df = pd.DataFrame(all_features)

        # 动态调整列顺序
        available_cols = ['文件编号', '方向', '样本编号', '时间点']
        for feature_group in [
            extractor.time_domain_feature_names,
            extractor.freq_domain_feature_names,
            extractor.time_freq_feature_names,
            extractor.wavelet_feature_names,
            extractor.wavelet_band_names
        ]:
            available_cols.extend([col for col in feature_group if col in features_df.columns])

        return features_df[available_cols]
    else:
        return pd.DataFrame()


def process_multiple_files(input_dir: str, output_path: str, file_range: Tuple[int, int],
                           sampling_freq: int = 25600, window_sec: float = 1.28,
                           overlap_ratio: float = 0.0) -> None:
    """
    批量处理多个CSV文件并保存结果

    参数:
        input_dir: 输入文件夹路径
        output_path: 输出CSV文件路径
        file_range: 文件编号范围，如(1, 123)表示处理文件1到123
    """
    all_dfs = []

    # 遍历文件范围
    for i in tqdm(range(file_range[0], file_range[1] + 1), desc="处理文件"):
        file_path = os.path.join(input_dir, f"{i}.csv")

        if not os.path.exists(file_path):
            print(f"文件不存在: {file_path}")
            continue

        # 处理单个文件
        df = process_single_file(
            file_path=file_path,
            sampling_freq=sampling_freq,
            window_sec=window_sec,
            overlap_ratio=overlap_ratio
        )

        if not df.empty:
            all_dfs.append(df)

    # 合并所有DataFrame
    if all_dfs:
        combined_df = pd.concat(all_dfs, ignore_index=True)

        # 保存结果
        combined_df.to_csv(output_path, index=False, encoding='utf-8-sig')
        print(f"特征提取完成，结果已保存至 {output_path}")
        print(f"共处理 {len(all_dfs)} 个文件，提取了 {len(combined_df)} 个特征样本")
    else:
        print("未提取到任何特征，请检查文件路径和格式")


input_directory = r"D:\轴承寿命预测\Bearing1_1"  # 包含1-123.csv的文件夹
output_file = "./combined_features.csv"  # 输出特征文件

process_multiple_files(
    input_dir=input_directory,
    output_path=output_file,
    file_range=(1, 123),  # 处理文件1到123
    sampling_freq=25600,  # 采样频率
    window_sec=1.28,  # 窗口大小(秒)
    overlap_ratio=0.0  # 无重叠
)